import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文
plt.rcParams['axes.unicode_minus'] = False    # 正常显示负号

# 实验二：有理函数和部分分式分解 (简化版)
print("实验二：有理函数和部分分式分解")
print("="*40)

# 实现拉格朗日插值法（基于习题6、7）
def lagrange_interpolation(z_vals, f_vals, z):
    n = len(z_vals)
    result = 0.0
    for i in range(n):
        basis = 1.0
        for j in range(n):
            if i != j:
                basis *= (z - z_vals[j]) / (z_vals[i] - z_vals[j])
        result += f_vals[i] * basis
    return result

# 验证习题6的插值公式
print("1. 验证拉格朗日插值公式（习题6）:")
def f(z): return z**2 + 2*z + 1  # 次数<3的多项式

z_vals = np.array([1, 2, 3])
f_vals = f(z_vals)

print("给定3个点:")
for i in range(len(z_vals)):
    print(f"  z{i+1} = {z_vals[i]}, f(z{i+1}) = {f_vals[i]}")

# 验证插值
test_point = 1.5
interpolated_value = lagrange_interpolation(z_vals, f_vals, test_point)
actual_value = f(test_point)
print(f"\n在 z = {test_point} 处:")
print(f"  插值结果: {interpolated_value:.6f}")
print(f"  实际值: {actual_value:.6f}")
print(f"  误差: {abs(interpolated_value - actual_value):.2e}")

# 部分分式分解（习题11）
print("\n2. 部分分式分解 f(z) = z^4/(z^3 - 1):")
A = 1/3  # 通过留数方法计算的系数
print(f"  z^4/(z^3 - 1) = z + {A:.4f}/(z-1) + (Bz + C)/(z^2 + z + 1)")

# 验证分解
z_test = 2 + 1j
original = z_test**4 / (z_test**3 - 1)
decomposed = z_test + A/(z_test - 1)
print(f"\n在 z = {z_test}:")
print(f"  原函数值: {original:.6f}")
print(f"  分解近似值: {decomposed:.6f}")
print(f"  误差: {abs(original - decomposed):.2e}")

# 可视化有理函数
print("\n3. 可视化有理函数 f(z) = z^4/(z^3 - 1):")
x = np.linspace(-2, 2, 40)
y = np.linspace(-2, 2, 40)
X, Y = np.meshgrid(x, y)
Z = X + 1j*Y

# 计算函数值（处理奇点）
F = np.where(abs(Z**3 - 1) > 1e-10, Z**4 / (Z**3 - 1), np.nan)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))

# 绘制模长
magnitude = np.abs(F)
im1 = ax1.imshow(magnitude, extent=[-2, 2, -2, 2], origin='lower', cmap='viridis')
ax1.set_title('|f(z)|')
plt.colorbar(im1, ax=ax1)

# 绘制相位
phase = np.angle(F)
im2 = ax2.imshow(phase, extent=[-2, 2, -2, 2], origin='lower', cmap='hsv')
ax2.set_title('arg(f(z))')
plt.colorbar(im2, ax=ax2)

plt.tight_layout()
plt.show()